import logging
from ipaddress import IPv4Address
from time import sleep
from typing import Optional, Set

import pymssql

from common.credentials import Credentials
from common.types import NetworkPort

logger = logging.getLogger(__name__)
QUERY_BUFFER_TIME = 0.5
CONNECTION_CHECK_QUERY = "SELECT 1"


class MSSQLClient:
    def __init__(self, server_timeout: float, query_buffer_time: float = QUERY_BUFFER_TIME):
        self._cursor: Optional[pymssql.Cursor] = None
        self._server_timeout = server_timeout
        self._query_buffer_time = query_buffer_time

    def login(self, ip: IPv4Address, ports: Set[NetworkPort], credentials: Credentials):
        """
        Login to MSSQL server using the given credentials

        :param ip: IP address of the MSSQL server
        :param ports: Set of ports to try to login with
        :param credentials: Credentials to use for login
        :raises RuntimeError: If login fails
        """
        for port in ports:
            if self._login(ip, port, credentials):
                return

        raise RuntimeError("Failed to authenticate with the MSSQL server with provided credentials")

    def _login(self, ip: IPv4Address, port: NetworkPort, credentials: Credentials) -> bool:
        try:
            connection = pymssql.connect(
                str(ip),
                credentials.identity.username,
                credentials.secret.password.get_secret_value(),
                port=port,
                login_timeout=self._server_timeout,
                timeout=self._server_timeout,
            )
            self._cursor = connection.cursor()
            self._cursor.execute(CONNECTION_CHECK_QUERY)
            self._cursor.fetchall()
            return True
        except pymssql.OperationalError:
            logger.debug("MSSQL connection is not active")
            return False

    def run_command(self, command: str):
        """
        Run a command on the MSSQL server

        :param command: The command to run
        """
        if self._cursor is None:
            raise RuntimeError("You are not logged in. Please log in before running commands.")
        try:
            self._cursor.execute(command)
            sleep(self._query_buffer_time)
        except pymssql.Error as err:
            raise RuntimeError("Failed to execute the command") from err
